package com.example.ai.embedding;

import com.example.ai.etl.MyTokenTextSplitter;
import com.example.ai.etl.reader.MyJsonReader;
import com.example.ai.etl.reader.MyTikaDocumentReader;
import io.micrometer.observation.ObservationRegistry;
import io.qdrant.client.QdrantClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.ai.vectorstore.qdrant.autoconfigure.QdrantVectorStoreProperties;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.util.List;

@RestController
@RequestMapping("/verctor/qdrant")
public class QdrantController {

    @Autowired
    EmbeddingModel embeddingModel;
    @Autowired
    QdrantVectorStoreProperties properties;
    @Autowired
    QdrantClient qdrantClient;
    @Autowired
    ObjectProvider<ObservationRegistry> observationRegistry;
    @Autowired
    ObjectProvider<VectorStoreObservationConvention> customObservationConvention;
    @Autowired
    BatchingStrategy batchingStrategy;

    @Autowired
    MyJsonReader myJsonReader;

    @Autowired
    MyTikaDocumentReader tikaDocumentReader;
    @Autowired
    MyTokenTextSplitter tokenTextSplitter;

    @PostMapping("/init")
    public ResponseEntity<?> initData(@RequestBody QdrantETLReq req) {
        VectorStore store = QdrantVectorStore.builder(qdrantClient, embeddingModel)
                .collectionName(req.getCollectionName())
                .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
                .customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
                .batchingStrategy(batchingStrategy)
                .build();

        if(req.getClassFilePath().endsWith("json")){
            MyJsonMetadataGenerator myJsonMetadataGenerator = new MyJsonMetadataGenerator();
            List<Document> documents = myJsonReader.get(req.getClassFilePath(),myJsonMetadataGenerator);
            store.add(documents);
        }

        if(req.getClassFilePath().endsWith("docx")){
            //"labour_dispute";
            List<Document> documents = tikaDocumentReader.loadText(req.getClassFilePath());
            List<Document> splitedDocs = tokenTextSplitter.splitDocuments(documents);
            store.add(splitedDocs);
        }

        return ResponseEntity.ok("success");
    }

    @PostMapping("/similarSearch")
    public ResponseEntity<?> similarSearch(@RequestBody  QdrantSimilarSearchReq req) {
        VectorStore store = QdrantVectorStore.builder(qdrantClient, embeddingModel)
                .collectionName(req.getCollectionName())
                .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
                .customObservationConvention(customObservationConvention.getIfAvailable(() -> null))
                .batchingStrategy(batchingStrategy)
                .build();
        List<Document> results = store.similaritySearch(SearchRequest.builder()
                .query(req.getContent())
                .topK(req.getTopK())
                .similarityThreshold(req.getThreshold())
                .filterExpression(req.getFilter())
                .build());
        return ResponseEntity.ok(results);
    }
}
